package com.alex.aichattongyiqianwen.controller;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import com.fasterxml.jackson.annotation.JsonProperty;
import io.reactivex.Flowable;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.atomic.AtomicLong;

@RestController
@RequestMapping("/api/chat")
public class ChatController {
    // 用于生成唯一 requestId
    private static final AtomicLong COUNTER = new AtomicLong(1);

    static class ChatRequest {
        public List<Map<String, String>> messages;
        public boolean deepthinking;
        public String preset;

        @JsonProperty("enable_search")
        public Boolean enable_search;

        @JsonProperty("enable_thinking")
        public Boolean enable_thinking;
    }

    @PostMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public void streamChat(@RequestBody ChatRequest req, HttpServletResponse response)
            throws IOException {

        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        response.setCharacterEncoding("UTF-8");

        String requestId = UUID.randomUUID().toString();

        // 构建消息列表
        List<Message> messageList = new ArrayList<>();
        if (req.preset != null && !req.preset.trim().isEmpty()) {
            messageList.add(Message.builder()
                    .role(Role.SYSTEM.getValue())
                    .content(req.preset)
                    .build());
        }
        if (req.messages != null) {
            for (Map<String, String> m : req.messages) {
                String role = m.getOrDefault("role", "user");
                String content = m.getOrDefault("content", "");
                messageList.add(Message.builder()
                        .role(role)
                        .content(content)
                        .build());
            }
        }

        // extraBody 参数配置
        Map<String, Object> extraBody = Map.of(
                "stream", true,
                "enable_search", req.enable_search != null ? req.enable_search : true,
                "enable_thinking", req.enable_thinking != null ? req.enable_thinking : true
        );

        GenerationParam param = GenerationParam.builder()
                .apiKey(System.getenv("bailian-api"))
                .model("qwen-plus-2025-04-28")
                .messages(messageList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .parameters(extraBody) // extraBody要包含 enable_thinking:true
                .incrementalOutput(true) // 这一行很关键！！！
                .build();

        Generation gen = new Generation();

        try {
            Flowable<?> stream = gen.streamCall(param);

            stream.blockingForEach(chunk -> {
                try {
                    // 必须转成 JSON
                    String chunkJson = JsonUtils.toJson(chunk);

                    // 你可以区分 reasoning 和 reply 类型
                    String eventType = "reply";
                    if (chunkJson.contains("\"reasoning\"")) {
                        eventType = "reasoning";
                    }

                    // 组装 SSE 消息
                    StringBuilder sb = new StringBuilder();
                    sb.append("event: ").append(eventType).append("\n");
                    sb.append("data: ").append(chunkJson.replace("\n", "")).append("\n\n");

                    response.getWriter().write(sb.toString());
                    response.getWriter().flush();
                } catch (IOException e) {
                    // 客户端断开等
                }
            });
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            // 输出错误 SSE
            String errMsg = e.getMessage();
            String errJson = "{\"error\": \"" + errMsg.replace("\"", "'") + "\"}";
            response.getWriter().write("event: error\ndata: " + errJson + "\n\n");
            response.getWriter().flush();
        }
    }
}